import math
import cv2 as cv
import numpy as np
from scipy import signal
import os
from superposition import superPosition, superPositionAlpha

'''
函数描述：去除图片中的噪点。
JZhou@20211124
'''

def pascal_smooth(n):
  	# 返回 n 阶的非归一化的高斯平滑算子
    pascal_smooth = np.zeros([1, n], np.float32)
    for i in range(n):
        pascal_smooth[0][i] = math.factorial(n-1) / (math.factorial(i) * math.factorial(n-1-i))
    return pascal_smooth

def pascal_diff(n):
  	# 返回 n 阶的差分算子
    pascal_diff = np.zeros([1, n], np.float32)
    pascal_smooth_previous = pascal_smooth(n-1)
    for i in range(n):
        if i == 0:
            # 恒等于 1
            pascal_diff[0][i] = pascal_smooth_previous[0][i]
        elif i == n-1:
            # 恒等于 -1
            pascal_diff[0][i] = - pascal_smooth_previous[0][i-1]
        else:
            pascal_diff[0][i] = pascal_smooth_previous[0][i] - pascal_smooth_previous[0][i-1]
    return pascal_diff


def get_sobel_kernel(n):
    pascal_smooth_kernel = pascal_smooth(n)
    pascal_diff_kernel = pascal_diff(n)
    # 水平方向的卷积核
    sobel_kerenl_x = signal.convolve2d(pascal_smooth_kernel.transpose(), pascal_diff_kernel, mode='full')
    # 垂直方向的卷积核
    sobel_kerenl_y = signal.convolve2d(pascal_smooth_kernel, pascal_diff_kernel.transpose(), mode='full')
    return sobel_kerenl_x, sobel_kerenl_y

def sobel(img, n):
    rows, cols = img.shape[:2]
    # 平滑算子
    pascal_smooth_kernel = pascal_smooth(n)
    # 差分算子
    pascal_diff_kernel = pascal_diff(n)
    # 水平方向上的 sobel 核卷积
    # 先进行垂直方向的平滑
    img_sobel_x = signal.convolve2d(img, pascal_smooth_kernel.transpose(), mode='same')
    # 再进行水平方向上的差分
    img_sobel_x = signal.convolve2d(img_sobel_x, pascal_diff_kernel, mode='same')
    # 垂直方向上的 sobel 核卷积
    img_sobel_y = signal.convolve2d(img, pascal_smooth_kernel, mode='same')
    img_sobel_y = signal.convolve2d(img_sobel_y, pascal_diff_kernel.transpose(), mode='same')

    return img_sobel_x, img_sobel_y


def colorMask(img):
    img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    hsv_img = cv.cvtColor(img, cv.COLOR_RGB2HSV)
    light_orange = (11,43,46)
    dark_orange = (25,255,255)

    light_white = (0,0,0)
    dark_white = (180,255,46)

    orange_mask = cv.inRange(hsv_img,light_orange,dark_orange)
    white_mask = cv.inRange(hsv_img,light_white,dark_white)
    mask = orange_mask+white_mask
    # mask = white_mask
    # cv.imwrite('/Users/vine/Desktop/2mask.jpg', mask)

    return mask

def colorMask2(img):
    #转变为HSV颜色空间
    # img_hsv=cv.cvtColor(img,cv.COLOR_BGR2HSV)
    #返回黄色区域的二值图像
    # img_yellow=cv2.inRange(img_original,(100,43,46),(125,255,255)) # 蓝色
    img_yellow=cv.inRange(img,(0,0,0),(180,255,46)) # 黑色
    # img_yellow=cv2.inRange(img_original,(0,43,46),(10,255,255)) # 红色
    # img_yellow=cv2.inRange(img_original,(0,43,46),(34,255,255)) # 红色
    ##输入图像与输入图像在掩模条件下按位与，得到掩模范围内的原图像
    img_specifiedColor=cv.bitwise_and(img,img,mask=img_yellow)
    # cv.imwrite('/Users/vine/Desktop/colorMask11.jpg', img_specifiedColor)
    return img_specifiedColor

def handlePic(img, save_path):
    imgInitial = img.copy()
    blurred = cv.pyrMeanShiftFiltering(img, 6, 30) # 均值迁移去噪声
    cv.imwrite('/Users/vine/Desktop/blurred.jpg', blurred)
    # 颜色掩膜
    img = colorMask2(img)
    # 二值化
    img_gray = cv.cvtColor(blurred,cv.COLOR_RGB2GRAY) # 将图片转为灰度图
    blur = cv.GaussianBlur(img_gray, (5, 5), 0)
    cv.imwrite('/Users/vine/Desktop/blur.jpg', blur)
    ret3, img = cv.threshold(blur, 40, 255, cv.THRESH_BINARY)

    img_sobel_x, img_sobel_y = sobel(img, 3)
    img_sobel_x_c, img_sobel_y_c = img_sobel_x.copy(), img_sobel_y.copy()
    img_sobel_x_c, img_sobel_y_c = abs(img_sobel_x_c), abs(img_sobel_y_c)
    img_sobel_x_c[img_sobel_x_c>255] = 255
    img_sobel_y_c[img_sobel_y_c>255] = 255
    img_sobel_x_c = img_sobel_x_c.astype(np.uint8)
    img_sobel_y_c = img_sobel_y_c.astype(np.uint8)
    # 平方和开方的方式
    edge = np.sqrt(np.power(img_sobel_x, 2.0) + np.power(img_sobel_y, 2.0))

    # 直接截断显示
    edge_c = edge.copy()
    edge_c[edge_c > 255] = 255
    edge_c = edge_c.astype(np.uint8)
    cv.imwrite("/Users/vine/Desktop/edge_c.jpg", edge_c)

    # 在原图中标明边缘线
    # img_processing = superPosition(edge_c, imgInitial)
    img_processing = superPositionAlpha(edge_c, imgInitial)

    # 这里改一下名称，最后三位强制改为jpg
    if save_path[-3:] == "png":
        save_path=save_path.replace("png", "jpg")
    cv.imwrite(save_path, img_processing)
    print("Save as : " + save_path)

def show_files(path, all_files):
    # 首先遍历当前目录所有文件及文件夹
    file_list = os.listdir(path)
    # 准备循环判断每个元素是否是文件夹还是文件，是文件的话，把名称传入list，是文件夹的话，递归
    for file in file_list:
        # 利用os.path.join()方法取得路径全名，并存入cur_path变量，否则每次只能遍历一层目录
        cur_path = os.path.join(path, file)
        # 判断是否是文件夹
        if os.path.isdir(cur_path):
            show_files(cur_path, all_files)
        else:
            # 拼接文件路径
            all_files.append(path+"/"+file)
    return all_files

if __name__ == '__main__':
    # test_file_path = '/Users/vine/Documents/Winglab/ISPTestCode/testPic/Clip'
    # if not os.path.exists("/Users/vine/Documents/Winglab/ISPTestCode/resPic/Clip"):
    #     os.mkdir("/Users/vine/Documents/Winglab/ISPTestCode/resPic/Clip")
    # res_file_path = '/Users/vine/Documents/Winglab/ISPTestCode/resPic/Clip'

    # # 首先遍历文件夹，然后对每个文件进行处理
    # # 传入空的list接收文件名
    # contents = show_files(test_file_path, [])
    # # 循环打印show_files函数返回的文件名列表
    # for content in contents:
    #     # print(content)
    #     # 判断是否为图片
    #     if content.endswith('jpg') or content.endswith('png'):
    #         # print("processing : "+content)
    #         # print(res_file_path + "/" +os.path.basename(content))
    #         img = cv.imread(content)
    #         handlePic(img,res_file_path + "/" +os.path.basename(content))

    # 单张图片
    res_file_path = "/Users/vine/Desktop"
    path = "/Users/vine/Documents/Winglab/ISPTestCode/testPic/Clip_V1/2021091120354001191G-CL0615-6.jpg"
    img = cv.imread(path)
    handlePic(img,res_file_path + "/" +os.path.basename(path))
